streamlit-nightly 1.45.1.dev20250511__py3-none-any.whl → 1.45.2.dev20250513__py3-none-any.whl

streamlit/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 # isort: skip_file
+# ruff: noqa: E402, PLC0414, A001
 
 """Streamlit.
 

streamlit/auth_util.py

@@ -32,21 +32,21 @@ if TYPE_CHECKING:
 class AuthCache:
     """Simple cache implementation for storing info required for Authlib."""
 
-    def __init__(self):
-        self.cache = {}
+    def __init__(self) -> None:
+        self.cache: dict[str, Any] = {}
 
-    def get(self, key):
+    def get(self, key: str) -> Any:
         return self.cache.get(key)
 
     # for set method, we are follow the same signature used in Authlib
     # the expires_in is not used in our case
-    def set(self, key, value, expires_in):  # noqa: ARG002
+    def set(self, key: str, value: Any, expires_in: int | None = None) -> None:  # noqa: ARG002
         self.cache[key] = value
 
-    def get_dict(self):
+    def get_dict(self) -> dict[str, Any]:
         return self.cache
 
-    def delete(self, key):
+    def delete(self, key: str) -> None:
         self.cache.pop(key, None)
 
 

streamlit/commands/echo.py

@@ -19,12 +19,12 @@ import contextlib
 import re
 import textwrap
 import traceback
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
 
 from streamlit.runtime.metrics_util import gather_metrics
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable
+    from collections.abc import Generator, Iterable
 
 _SPACES_RE = re.compile("\\s*")
 _EMPTY_LINE_RE = re.compile("\\s*\n")
@@ -32,7 +32,9 @@ _EMPTY_LINE_RE = re.compile("\\s*\n")
 
 @gather_metrics("echo")
 @contextlib.contextmanager
-def echo(code_location="above"):
+def echo(
+    code_location: Literal["above", "below"] = "above",
+) -> Generator[None, None, None]:
     """Use in a `with` block to draw some code on the app, then execute it.
 
     Parameters

streamlit/commands/experimental_query_params.py

@@ -123,7 +123,7 @@ def _extract_key_query_params(
         item.lower()
         for sublist in [
             [value.lower() for value in query_params[key]]
-            for key in query_params.keys()
+            for key in query_params
             if key.lower() == param_key and query_params.get(key)
         ]
         for item in sublist

streamlit/commands/navigation.py

@@ -311,10 +311,7 @@ def _navigation(
     # Build the pagehash-to-pageinfo mapping.
     for section_header in nav_sections:
         for page in nav_sections[section_header]:
-            if isinstance(page._page, Path):
-                script_path = str(page._page)
-            else:
-                script_path = ""
+            script_path = str(page._page) if isinstance(page._page, Path) else ""
 
             script_hash = page._script_hash
             if script_hash in pagehash_to_pageinfo:

streamlit/components/v1/custom_component.py

@@ -183,7 +183,7 @@ And if you're using Streamlit Cloud, add "pyarrow" to your requirements.txt."""
             # other arguments change, and the component's iframe won't be
             # remounted on the frontend.
 
-            def marshall_element_args():
+            def marshall_element_args() -> None:
                 element.component_instance.json_args = serialized_json_args
                 element.component_instance.special_args.extend(special_args)
 
@@ -210,7 +210,7 @@ And if you're using Streamlit Cloud, add "pyarrow" to your requirements.txt."""
                 )
             element.component_instance.id = computed_id
 
-            def deserialize_component(ui_value):
+            def deserialize_component(ui_value: Any) -> Any:
                 # ui_value is an object from json, an ArrowTable proto, or a bytearray
                 return ui_value
 

streamlit/config.py

@@ -728,15 +728,12 @@ def _server_headless() -> bool:
     Default: false unless (1) we are on a Linux box where DISPLAY is unset, or
     (2) we are running in the Streamlit Atom plugin.
     """
-    if (
+    # Check if we are running in Linux and DISPLAY is unset
+    return (
         env_util.IS_LINUX_OR_BSD
         and not os.getenv("DISPLAY")
         and not os.getenv("WAYLAND_DISPLAY")
-    ):
-        # We're running in Linux and DISPLAY is unset
-        return True
-
-    return False
+    )
 
 
 _create_option(
@@ -1629,10 +1626,11 @@ def _check_conflicts() -> None:
         )
 
     # XSRF conflicts
-    if get_option("server.enableXsrfProtection"):
-        if not get_option("server.enableCORS") or get_option("global.developmentMode"):
-            LOGGER.warning(
-                """
+    if get_option("server.enableXsrfProtection") and (
+        not get_option("server.enableCORS") or get_option("global.developmentMode")
+    ):
+        LOGGER.warning(
+            """
 Warning: the config option 'server.enableCORS=false' is not compatible with
 'server.enableXsrfProtection=true'.
 As a result, 'server.enableCORS' is being overridden to 'true'.
@@ -1644,7 +1642,7 @@ cross-origin resource sharing.
 
 If cross origin resource sharing is required, please disable server.enableXsrfProtection.
             """
-            )
+        )
 
 
 def _set_development_mode() -> None:
@@ -1653,7 +1651,7 @@ def _set_development_mode() -> None:
 
 def on_config_parsed(
     func: Callable[[], None], force_connect: bool = False, lock: bool = False
-) -> Callable[[], bool]:
+) -> Callable[[], None]:
     """Wait for the config file to be parsed then call func.
 
     If the config file has already been parsed, just calls func immediately
@@ -1673,7 +1671,7 @@ def on_config_parsed(
 
     Returns
     -------
-    Callable[[], bool]
+    Callable[[], None]
         A function that the caller can use to deregister func.
     """
 
@@ -1682,13 +1680,13 @@ def on_config_parsed(
     # leading to a memory leak because the Signal will keep a reference of the
     # callable argument. When the callable argument is an object method, then
     # the reference to that object won't be released.
-    def receiver(_):
-        return func_with_lock()
+    def receiver(_: Any) -> None:
+        func_with_lock()
 
-    def disconnect():
-        return _on_config_parsed.disconnect(receiver)
+    def disconnect() -> None:
+        _on_config_parsed.disconnect(receiver)
 
-    def func_with_lock():
+    def func_with_lock() -> None:
         if lock:
             with _config_lock:
                 func()

streamlit/config_util.py

@@ -54,19 +54,19 @@ def show_config(
         )
     )
 
-    def append_desc(text):
+    def append_desc(text: str) -> None:
         out.append("# " + cli_util.style_for_cli(text, bold=True))
 
-    def append_comment(text):
+    def append_comment(text: str) -> None:
         out.append("# " + cli_util.style_for_cli(text))
 
-    def append_section(text):
+    def append_section(text: str) -> None:
         out.append(cli_util.style_for_cli(text, bold=True, fg="green"))
 
-    def append_setting(text):
+    def append_setting(text: str) -> None:
         out.append(cli_util.style_for_cli(text, fg="green"))
 
-    for section in section_descriptions.keys():
+    for section in section_descriptions:
         # We inject a fake config section used for unit tests that we exclude here as
         # its options are often missing required properties, which confuses the code
         # below.
@@ -178,8 +178,7 @@ def _clean_paragraphs(txt: str) -> list[str]:
     # Strip both leading and trailing newlines.
     txt = txt.strip("\n")
     paragraphs = txt.split("\n\n")
-    cleaned_paragraphs = [
+    return [
         "\n".join(_clean(line) for line in paragraph.split("\n"))
         for paragraph in paragraphs
     ]
-    return cleaned_paragraphs

streamlit/connections/sql_connection.py

@@ -207,7 +207,7 @@ class SQLConnection(BaseConnection["Engine"]):
                 host=conn_params["host"],
                 port=int(conn_params["port"]) if "port" in conn_params else None,
                 database=conn_params.get("database"),
-                query=conn_params["query"] if "query" in conn_params else None,
+                query=conn_params.get("query", {}),
             )
 
         create_engine_kwargs = ChainMap(

streamlit/dataframe_util.py

@@ -1067,7 +1067,7 @@ def is_colum_type_arrow_incompatible(column: Series[Any] | Index) -> bool:
             "complex",
         ]:
             return True
-        elif inferred_type == "mixed":
+        if inferred_type == "mixed":
             # This includes most of the more complex/custom types (objects, dicts,
             # lists, ...)
             if len(column) == 0 or not hasattr(column, "iloc"):
@@ -1079,7 +1079,7 @@ def is_colum_type_arrow_incompatible(column: Series[Any] | Index) -> bool:
             # Get the first value to check if it is a supported list-like type.
             first_value = column.iloc[0]
 
-            if (
+            if (  # noqa: SIM103
                 not is_list_like(first_value)
                 # dicts are list-like, but have issues in Arrow JS (see comments in
                 # Quiver.ts)
@@ -1164,53 +1164,53 @@ def determine_data_format(input_data: Any) -> DataFormat:
 
     if input_data is None:
         return DataFormat.EMPTY
-    elif isinstance(input_data, pd.DataFrame):
+    if isinstance(input_data, pd.DataFrame):
         return DataFormat.PANDAS_DATAFRAME
-    elif isinstance(input_data, np.ndarray):
+    if isinstance(input_data, np.ndarray):
         if len(cast("NumpyShape", input_data.shape)) == 1:
             # For technical reasons, we need to distinguish one
             # one-dimensional numpy array from multidimensional ones.
             return DataFormat.NUMPY_LIST
         return DataFormat.NUMPY_MATRIX
-    elif isinstance(input_data, pa.Table):
+    if isinstance(input_data, pa.Table):
         return DataFormat.PYARROW_TABLE
-    elif isinstance(input_data, pa.Array):
+    if isinstance(input_data, pa.Array):
         return DataFormat.PYARROW_ARRAY
-    elif isinstance(input_data, pd.Series):
+    if isinstance(input_data, pd.Series):
         return DataFormat.PANDAS_SERIES
-    elif isinstance(input_data, pd.Index):
+    if isinstance(input_data, pd.Index):
         return DataFormat.PANDAS_INDEX
-    elif is_pandas_styler(input_data):
+    if is_pandas_styler(input_data):
         return DataFormat.PANDAS_STYLER
-    elif isinstance(input_data, pd.api.extensions.ExtensionArray):
+    if isinstance(input_data, pd.api.extensions.ExtensionArray):
         return DataFormat.PANDAS_ARRAY
-    elif is_polars_series(input_data):
+    if is_polars_series(input_data):
         return DataFormat.POLARS_SERIES
-    elif is_polars_dataframe(input_data):
+    if is_polars_dataframe(input_data):
         return DataFormat.POLARS_DATAFRAME
-    elif is_polars_lazyframe(input_data):
+    if is_polars_lazyframe(input_data):
         return DataFormat.POLARS_LAZYFRAME
-    elif is_modin_data_object(input_data):
+    if is_modin_data_object(input_data):
         return DataFormat.MODIN_OBJECT
-    elif is_snowpandas_data_object(input_data):
+    if is_snowpandas_data_object(input_data):
         return DataFormat.SNOWPANDAS_OBJECT
-    elif is_pyspark_data_object(input_data):
+    if is_pyspark_data_object(input_data):
         return DataFormat.PYSPARK_OBJECT
-    elif is_xarray_dataset(input_data):
+    if is_xarray_dataset(input_data):
         return DataFormat.XARRAY_DATASET
-    elif is_xarray_data_array(input_data):
+    if is_xarray_data_array(input_data):
         return DataFormat.XARRAY_DATA_ARRAY
-    elif is_ray_dataset(input_data):
+    if is_ray_dataset(input_data):
         return DataFormat.RAY_DATASET
-    elif is_dask_object(input_data):
+    if is_dask_object(input_data):
         return DataFormat.DASK_OBJECT
-    elif is_snowpark_data_object(input_data) or is_snowpark_row_list(input_data):
+    if is_snowpark_data_object(input_data) or is_snowpark_row_list(input_data):
         return DataFormat.SNOWPARK_OBJECT
-    elif is_duckdb_relation(input_data):
+    if is_duckdb_relation(input_data):
         return DataFormat.DUCKDB_RELATION
-    elif is_dbapi_cursor(input_data):
+    if is_dbapi_cursor(input_data):
         return DataFormat.DBAPI_CURSOR
-    elif (
+    if (
         isinstance(
             input_data,
             (ChainMap, UserDict, MappingProxyType),
@@ -1221,9 +1221,9 @@ def determine_data_format(input_data: Any) -> DataFormat:
         or is_pydantic_model(input_data)
     ):
         return DataFormat.KEY_VALUE_DICT
-    elif isinstance(input_data, (ItemsView, enumerate)):
+    if isinstance(input_data, (ItemsView, enumerate)):
         return DataFormat.LIST_OF_ROWS
-    elif isinstance(input_data, (list, tuple, set, frozenset)):
+    if isinstance(input_data, (list, tuple, set, frozenset)):
         if _is_list_of_scalars(input_data):
             # -> one-dimensional data structure
             if isinstance(input_data, tuple):
@@ -1231,15 +1231,14 @@ def determine_data_format(input_data: Any) -> DataFormat:
             if isinstance(input_data, (set, frozenset)):
                 return DataFormat.SET_OF_VALUES
             return DataFormat.LIST_OF_VALUES
-        else:
-            # -> Multi-dimensional data structure
-            # This should always contain at least one element,
-            # otherwise the values type from infer_dtype would have been empty
-            first_element = next(iter(input_data))
-            if isinstance(first_element, dict):
-                return DataFormat.LIST_OF_RECORDS
-            if isinstance(first_element, (list, tuple, set, frozenset)):
-                return DataFormat.LIST_OF_ROWS
+        # -> Multi-dimensional data structure
+        # This should always contain at least one element,
+        # otherwise the values type from infer_dtype would have been empty
+        first_element = next(iter(input_data))
+        if isinstance(first_element, dict):
+            return DataFormat.LIST_OF_RECORDS
+        if isinstance(first_element, (list, tuple, set, frozenset)):
+            return DataFormat.LIST_OF_ROWS
     elif isinstance(input_data, (dict, Mapping)):
         if not input_data:
             return DataFormat.KEY_VALUE_DICT
@@ -1340,7 +1339,7 @@ def convert_pandas_df_to_data_format(
         DataFormat.SNOWPARK_OBJECT,
     }:
         return df
-    elif data_format == DataFormat.NUMPY_LIST:
+    if data_format == DataFormat.NUMPY_LIST:
         import numpy as np
 
         # It's a 1-dimensional array, so we only return
@@ -1348,48 +1347,48 @@ def convert_pandas_df_to_data_format(
         # Calling to_numpy() on the full DataFrame would result in:
         # [[1], [2]] instead of [1, 2]
         return np.ndarray(0) if df.empty else df.iloc[:, 0].to_numpy()
-    elif data_format == DataFormat.NUMPY_MATRIX:
+    if data_format == DataFormat.NUMPY_MATRIX:
         import numpy as np
 
         return np.ndarray(0) if df.empty else df.to_numpy()
-    elif data_format == DataFormat.PYARROW_TABLE:
+    if data_format == DataFormat.PYARROW_TABLE:
         import pyarrow as pa
 
         return pa.Table.from_pandas(df)
-    elif data_format == DataFormat.PYARROW_ARRAY:
+    if data_format == DataFormat.PYARROW_ARRAY:
         import pyarrow as pa
 
         return pa.Array.from_pandas(_pandas_df_to_series(df))
-    elif data_format == DataFormat.PANDAS_SERIES:
+    if data_format == DataFormat.PANDAS_SERIES:
         return _pandas_df_to_series(df)
-    elif data_format in {DataFormat.POLARS_DATAFRAME, DataFormat.POLARS_LAZYFRAME}:
+    if data_format in {DataFormat.POLARS_DATAFRAME, DataFormat.POLARS_LAZYFRAME}:
         import polars as pl  # type: ignore[import-not-found]
 
         return pl.from_pandas(df)
-    elif data_format == DataFormat.POLARS_SERIES:
+    if data_format == DataFormat.POLARS_SERIES:
         import polars as pl
 
         return pl.from_pandas(_pandas_df_to_series(df))
-    elif data_format == DataFormat.XARRAY_DATASET:
+    if data_format == DataFormat.XARRAY_DATASET:
         import xarray as xr  # type: ignore[import-not-found]
 
         return xr.Dataset.from_dataframe(df)
-    elif data_format == DataFormat.XARRAY_DATA_ARRAY:
+    if data_format == DataFormat.XARRAY_DATA_ARRAY:
         import xarray as xr
 
         return xr.DataArray.from_series(_pandas_df_to_series(df))
-    elif data_format == DataFormat.LIST_OF_RECORDS:
+    if data_format == DataFormat.LIST_OF_RECORDS:
         return _unify_missing_values(df).to_dict(orient="records")
-    elif data_format == DataFormat.LIST_OF_ROWS:
+    if data_format == DataFormat.LIST_OF_ROWS:
         # to_numpy converts the dataframe to a list of rows
         return _unify_missing_values(df).to_numpy().tolist()
-    elif data_format == DataFormat.COLUMN_INDEX_MAPPING:
+    if data_format == DataFormat.COLUMN_INDEX_MAPPING:
         return _unify_missing_values(df).to_dict(orient="dict")
-    elif data_format == DataFormat.COLUMN_VALUE_MAPPING:
+    if data_format == DataFormat.COLUMN_VALUE_MAPPING:
         return _unify_missing_values(df).to_dict(orient="list")
-    elif data_format == DataFormat.COLUMN_SERIES_MAPPING:
+    if data_format == DataFormat.COLUMN_SERIES_MAPPING:
         return df.to_dict(orient="series")
-    elif data_format in [
+    if data_format in [
         DataFormat.LIST_OF_VALUES,
         DataFormat.TUPLE_OF_VALUES,
         DataFormat.SET_OF_VALUES,
@@ -1409,7 +1408,7 @@ def convert_pandas_df_to_data_format(
         if data_format == DataFormat.SET_OF_VALUES:
             return set(return_list)
         return return_list
-    elif data_format == DataFormat.KEY_VALUE_DICT:
+    if data_format == DataFormat.KEY_VALUE_DICT:
         df = _unify_missing_values(df)
         # The key is expected to be the index -> this will return the first column
         # as a dict with index as key.

streamlit/delta_generator.py

@@ -522,9 +522,12 @@ class DeltaGenerator(
 
     def _block(
         self,
-        block_proto: Block_pb2.Block = Block_pb2.Block(),
+        block_proto: Block_pb2.Block | None = None,
         dg_type: type | None = None,
     ) -> DeltaGenerator:
+        if block_proto is None:
+            block_proto = Block_pb2.Block()
+
         # Operate on the active DeltaGenerator, in case we're in a `with` block.
         dg = self._active_dg
 

streamlit/deprecation_util.py

@@ -93,7 +93,7 @@ def deprecate_func_name(
     """
 
     @functools.wraps(func)
-    def wrapped_func(*args, **kwargs):
+    def wrapped_func(*args: Any, **kwargs: Any) -> Any:
         result = func(*args, **kwargs)
         show_deprecation_warning(
             make_deprecated_name_warning(
@@ -167,7 +167,7 @@ def _create_deprecated_obj_wrapper(obj: TObj, show_warning: Callable[[], Any]) -
             show_warning()
 
     class Wrapper:
-        def __init__(self):
+        def __init__(self) -> None:
             # Override all the Wrapped object's magic functions
             for name in Wrapper._get_magic_functions(obj.__class__):
                 setattr(
@@ -176,7 +176,7 @@ def _create_deprecated_obj_wrapper(obj: TObj, show_warning: Callable[[], Any]) -
                     property(self._make_magic_function_proxy(name)),
                 )
 
-        def __getattr__(self, attr):
+        def __getattr__(self, attr: str) -> Any:
             # We handle __getattr__ separately from our other magic
             # functions. The wrapped class may not actually implement it,
             # but we still need to implement it to call all its normal
@@ -199,8 +199,8 @@ def _create_deprecated_obj_wrapper(obj: TObj, show_warning: Callable[[], Any]) -
             ]
 
         @staticmethod
-        def _make_magic_function_proxy(name):
-            def proxy(self, *args):  # noqa: ARG001
+        def _make_magic_function_proxy(name: str) -> Callable[[Any], Any]:
+            def proxy(_self: Any, *args: Any) -> Any:
                 maybe_show_warning()
                 return getattr(obj, name)
 

streamlit/elements/arrow.py

@@ -563,7 +563,7 @@ class ArrowMixin:
         if use_container_width is None:
             # If use_container_width was not explicitly set by the user, we set
             # it to True if width was not set explicitly, and False otherwise.
-            use_container_width = True if width is None else False
+            use_container_width = width is None
 
         proto.use_container_width = use_container_width
 

streamlit/elements/deck_gl_json_chart.py

@@ -456,10 +456,7 @@ class PydeckMixin:
 
         ctx = get_script_run_ctx()
 
-        if pydeck_obj is None:
-            spec = json.dumps(EMPTY_MAP)
-        else:
-            spec = pydeck_obj.to_json()
+        spec = json.dumps(EMPTY_MAP) if pydeck_obj is None else pydeck_obj.to_json()
 
         pydeck_proto.json = spec
         pydeck_proto.use_container_width = use_container_width

streamlit/elements/doc_string.py

@@ -283,9 +283,7 @@ def _get_variable_name_from_code_str(code: str) -> str | None:
     ):
         # A common pattern is to add "," at the end of a magic command to make it print.
         # This removes that final ",", so it looks nicer.
-        code = code.removesuffix(",")
-
-        return code
+        return code.removesuffix(",")
 
     arg_node = _get_stcommand_arg(tree)
 
@@ -427,7 +425,7 @@ def _get_weight(value: Any) -> int:
     return 0
 
 
-def _get_value(obj, var_name):
+def _get_value(obj: object, var_name: str | None) -> str | None:
     obj_value = _get_human_readable_value(obj)
 
     if obj_value is not None:
@@ -449,10 +447,7 @@ def _get_value(obj, var_name):
     sig = _get_signature(name_obj) or ""
 
     if name:
-        if module:
-            obj_value = f"{module}.{name}{sig}"
-        else:
-            obj_value = f"{name}{sig}"
+        obj_value = f"{module}.{name}{sig}" if module else f"{name}{sig}"
 
     if obj_value == var_name:
         # No need to repeat the same info.
@@ -462,7 +457,7 @@ def _get_value(obj, var_name):
     return obj_value
 
 
-def _get_human_readable_value(value):
+def _get_human_readable_value(value: Any) -> str | None:
     if isinstance(value, Secrets):
         # Don't want to read secrets.toml because that will show a warning if there's no
         # secrets.toml file.
@@ -511,7 +506,7 @@ def _is_computed_property(obj: object, attr_name: str) -> bool:
     return False
 
 
-def _get_members(obj):
+def _get_members(obj: object) -> list[MemberProto]:
     members_for_sorting = []
 
     for attr_name in dir(obj):

streamlit/elements/exception.py

@@ -232,10 +232,9 @@ def _format_syntax_error_message(exception: SyntaxError) -> str:
 
     """
     if exception.text:
-        if exception.offset is not None:
-            caret_indent = " " * max(exception.offset - 1, 0)
-        else:
-            caret_indent = ""
+        caret_indent = (
+            " " * max(exception.offset - 1, 0) if exception.offset is not None else ""
+        )
 
         return (
             'File "%(filename)s", line %(lineno)s\n'
@@ -362,9 +361,8 @@ def _split_list(
     saw_split_point = False
 
     for item in orig_list:
-        if not saw_split_point:
-            if split_point(item):
-                saw_split_point = True
+        if not saw_split_point and split_point(item):
+            saw_split_point = True
 
         if saw_split_point:
             after.append(item)

streamlit/elements/layouts.py

@@ -184,7 +184,7 @@ class LayoutsMixin:
         self,
         spec: SpecType,
         *,
-        gap: Literal["small", "medium", "large"] = "small",
+        gap: Literal["small", "medium", "large", "none"] = "small",
         vertical_alignment: Literal["top", "center", "bottom"] = "top",
         border: bool = False,
     ) -> list[DeltaGenerator]:
@@ -213,7 +213,7 @@ class LayoutsMixin:
               Or ``[1, 2, 3]`` creates three columns where the second one is two times
               the width of the first one, and the third one is three times that width.
 
-        gap : "small", "medium", or "large"
+        gap : "small", "medium", "large", or "none"
             The size of the gap between the columns. The default is ``"small"``.
 
         vertical_alignment : "top", "center", or "bottom"
@@ -359,11 +359,12 @@ class LayoutsMixin:
                 "small": GapSize.SMALL,
                 "medium": GapSize.MEDIUM,
                 "large": GapSize.LARGE,
+                "none": GapSize.NONE,
             }
 
             if isinstance(gap, str):
                 gap_size = gap.lower()
-                valid_sizes = ["small", "medium", "large"]
+                valid_sizes = gap_mapping.keys()
 
                 if gap_size in valid_sizes:
                     return gap_mapping[gap_size]